import glob 
import os
import sys
import subprocess  
import shutil  
import re
import json

from transformers import AutoModel, AutoTokenizer
from natsort import natsorted
from langchain_community.chat_models import ChatZhipuAI
from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
from langchain.output_parsers import StructuredOutputParser, ResponseSchema
from langchain_core.prompts import PromptTemplate

from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.image.exceptions import UnrecognizedImageError  
from PIL import Image


#降低视频帧率为17,生成cut.mp4
def reduce_frame(input_video, framerate=17, output_video ='cut.mp4'):  
 
    # 构建ffmpeg命令  
    command = [  
        'ffmpeg',  
        '-i', input_video,  # 输入文件  
        '-r', str(framerate),    # 设置帧率  
        output_video        # 输出文件  
    ]  
    # 执行命令  
    try:  
        # 使用subprocess.run执行命令，捕获输出和错误  
        print("正在使用ffmpeg转换视频文件的帧率(default=17)。")
        result = subprocess.run(command, capture_output=True, text=True, check=True)  
    except subprocess.CalledProcessError as e:  
        # 如果ffmpeg命令执行失败，则打印错误信息  
        print(f"ffmpeg命令执行失败，错误：")  
        print(e.stderr)  

#对cut.mp4关键帧提取，生成关键帧目录key_frames
def extract_i_frames(video_path, timestamp_file="timestamps.txt", output_dir="key_frames"):
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    
    # 获取I帧极其出现时间
    ffprobe_cmd = [
        'ffprobe',
        '-v', 'error',
        '-skip_frame', 'nokey',  # 跳过非关键帧
        '-select_streams', 'v:0',
        '-show_entries', 'frame=pts_time,pict_type',
        '-of', 'csv=print_section=0',
        video_path
    ]
    
    # 重筛I帧，并对时间信息进行处理
    output = subprocess.check_output(ffprobe_cmd).decode('utf-8')
    timestamps = []
    for line in output.splitlines():
        if line.endswith(',I'):
            raw_seconds = float(line.split(',')[0])
            
            minutes = int(raw_seconds // 60)
            seconds = int(raw_seconds % 60)
            formatted = f"{minutes}分{seconds}秒"

            timestamps.append(formatted) 
    
    # 保存已处理的时间戳信息
    with open(timestamp_file, 'w', encoding='utf-8') as f:
        f.writelines(timestamps)
    
    # 提取I帧图像
    ffmpeg_cmd = [
        'ffmpeg',
        '-i', video_path,
        '-vf', "select=eq(pict_type\\,I)",  # 选择I帧
        '-vsync', 'vfr',                   # 保持原始时间基准
        '-qscale:v', '4',                  # 控制图像质量（2-31，值越小质量越高）
        os.path.join(output_dir, '%04d.jpg')
    ]
    
    subprocess.run(ffmpeg_cmd)
    
    print(f"成功提取 {len(timestamps)} 个关键帧到 {output_dir} 目录")
    print(f"对应时间戳保存到 {timestamp_file} 目录")

    # 重命名文件为 'X分X秒_000X.jpg' 的格式
    frame_files = sorted(glob.glob(os.path.join(output_dir, '*.jpg')), 
                        key=lambda x: int(os.path.splitext(os.path.basename(x))[0]))

    for i, old_path in enumerate(frame_files):
        # 提取原文件编号，如'0001'
        original_number = os.path.splitext(os.path.basename(old_path))[0]
        # 确定时间戳或NaN
        if i < len(timestamps):
            ts = timestamps[i]  # 已经不带换行符
        else:
            ts = 'NaN'
        # 构建新文件名
        new_name = f"{ts}_{original_number}.jpg"
        new_path = os.path.join(output_dir, new_name)
        # 执行重命名
        os.rename(old_path, new_path)
    # print("重命名key_frame成功！")



def OCR_i_frames(image_dir='key_frames', output_txt='output.txt'):  
    print("正在下载OCR模型")

    tokenizer = AutoTokenizer.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True)
    model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
    model = model.eval().cuda()

    print("正在OCR关键帧内容")
    
    with open(output_txt, 'w', encoding='utf-8') as f:  
        cnt = 1
        for filename in natsorted(os.listdir(image_dir)):  
            if filename.lower().endswith(('.jpg', '.png')):  
                image_file = os.path.join(image_dir, filename)  
                
                try:
                    res = model.chat(tokenizer, image_file, ocr_type='ocr')
                    res = res.strip() or "<空识别结果>"
                except Exception as e:
                    res = f"识别失败: {str(e)}"
                
                # 增强版文件名解析
                parts = filename.rsplit('_', 1)
                if len(parts) == 2:
                    timestamp = parts[0].split('.')[0]  # 移除扩展名
                    timestamp = re.sub(r'[\\/:*?"<>|]', '', timestamp)
                else:
                    timestamp = "NaN"

                entry = f"{timestamp}: {res}\n\n"
                f.write(entry)
                cnt += 1
                
    print(f"\n所有{cnt-1}个关键帧内容已存放在{output_txt}")

#智谱轻言提取习题内容,  output.txt -> output.json
def AI_extract_questions(api_key, txt_file_path='output.txt', output_json_path='output.json'):
    print("正在使用大模型提取结构化习题内容")
    # 设置API密钥
    os.environ["ZHIPUAI_API_KEY"] = api_key

    # 初始化模型
    llm = ChatZhipuAI(
        model="glm-4-plus",
        temperature=0.5,
    )

    # 解析格式
    response_schemas = [
        ResponseSchema(name="questions", description="A list of extracted questions, each as a structured object", properties=
        [
            ResponseSchema(name="text", description="The text of the question"),
            ResponseSchema(name="type", description="The type of question "),
            ResponseSchema(name="answer", description="The correct answer to the question"),
            ResponseSchema(name="timestamp", description="The timestamp of the question")
        ])
    ]

    output_parser = StructuredOutputParser.from_response_schemas(response_schemas)
    format_instructions = output_parser.get_format_instructions()

    # prompt
    summarizing_prompt_template = """
    {format_instructions}

    请逐行分析文本，并提取文本中涉及到的地理、数学、物理、语文等习题。
    对于每道的习题(question)，输出格式应为json文件，json文件中必须包括问题内容(text)、类型(type)、答案(answer)和时间戳(timestamp)。
    其中问题的内容(text)、答案(answer)、时间戳(timestamp)在文本中都可以找到，但问题的类型(type)可能要由你来判断。
    如果内容(text)、类型(type)、答案(answer)和时间戳(timestamp)无法识别或者不确定，那么在json中用一个空格字符代替即可，但一定不可以没有其中任何一个。
    ---
    {content}
    """
    prompt = PromptTemplate.from_template(summarizing_prompt_template, partial_variables={'format_instructions': format_instructions})

    summarizing_chain = prompt | llm | output_parser
    # 读取文本文件
    def read_txt(file_path):
        with open(file_path, 'r', encoding='utf-8') as f:
            content = f.read()
        return content
    doc = read_txt(txt_file_path)

    # 提取问题并保存到JSON文件
    res = summarizing_chain.invoke({"content": doc})
    
    with open(output_json_path, 'w', encoding='utf-8') as f:
        json.dump(res, f, ensure_ascii=False)
    print("结构化习题内容已生成，保存在output.json")

# 将结构化习题json转为word
def json2word(json_file_path='output.json', images_dir='key_frames', output_docx_path='output.docx'):
    print("正在将结构化习题转化为word")
    doc = Document()
    
    # 添加异常处理确保JSON文件可读
    try:
        with open(json_file_path, 'r', encoding='utf-8') as f:
            json_data = json.load(f)
    except Exception as e:
        print(f"JSON文件读取失败: {str(e)}")
        return

    doc.add_heading('任务二：例题联系析取答案', 0)
    
    # 安全遍历questions列表
    for question in json_data.get('questions', []):  # 使用get避免KeyError
        # 使用get方法安全访问字段（关键修复）
        text = question.get('text', '')
        q_type = question.get('type', '')  # 修复KeyError的核心修改
        answer = question.get('answer', '')
        timestamp = question.get('timestamp', '')  # 同样处理timestamp字段

        # 修改后的条件判断（保持原有逻辑）
        if text:
            doc.add_heading(f"问题: {text}", level=1)
        if q_type:  # 现在安全访问
            doc.add_paragraph(f"类型: {q_type}")
        if answer:
            doc.add_paragraph(f"答案: {answer}")

        # 保持原有图片插入逻辑（仅修改字段访问方式）
        if timestamp:  # 使用已通过get获取的值
            doc.add_paragraph(f"时间戳: {timestamp}")  
            timestamps = timestamp.rstrip(':')  # 使用已处理过的timestamp变量
            current_dir = os.path.abspath('.')
            for i in range(1, 1000):  
                prefix = f"{i:04d}"  
                image_filename = f"{timestamps}_{prefix}.jpg"  
                image_path = os.path.join(current_dir, images_dir, image_filename)  
                
                if os.path.exists(image_path):  
                    try:  
                        doc.add_picture(image_path, width=Pt(400))  
                        break  
                    except UnrecognizedImageError:  
                        print(f"无法识别图片文件 {image_path}。跳过...")  
                        continue 
                elif i == 999:  
                    paragraph = doc.add_paragraph("对应的图片未找到。")  
                    paragraph.alignment = WD_ALIGN_PARAGRAPH.CENTER  
                else:  
                    continue  

    doc.save(output_docx_path)
    print("习题已生成!")
    return output_docx_path

def cleanup():
    resources = {
        'files': ['cut.mp4', 'output.json', 'output.txt', 'timestamps.txt'],
        'dirs': ['key_frames']
    }

    # 删除文件
    for file in resources['files']:
        path = os.path.abspath(file)  # 获取绝对路径
        if os.path.exists(path) and os.path.isfile(path):
            try:
                os.remove(path)
                print(f"✓ 已删除文件: {path}")
            except PermissionError:
                print(f"× 删除失败：文件被占用 {path}")
            except Exception as e:
                print(f"× 删除异常：{path} - {str(e)}")

    # 删除目录
    for dir in resources['dirs']:
        path = os.path.abspath(dir)
        if os.path.exists(path) and os.path.isdir(path):
            try:
                shutil.rmtree(path, ignore_errors=True)
                print(f"✓ 已删除目录: {path}")
            except Exception as e:
                print(f"× 目录删除失败：{path} - {str(e)}")
    
    #重命名output.docx为video_file_例题.
    # try:
    #     video_name = os.path.basename(video_path).split('.')[0]
    #     new_doc_name = f"{video_name}_例题.docx"
    #     os.rename("output.docx", new_doc_name)
    # except FileNotFoundError:
    #     print("output.docx 文件重命名失败")
    # except Exception as e:
    #     print(f"出现异常，output.docx重命名失败: {str(e)}")

if __name__ == "__main__":
    if len(sys.argv) < 2:  
        print("请提供一个视频文件作为参数,例如：python xx.py xx.mp4")  
    else:  
        video_file = sys.argv[1]  # 获取第一个命令行参数，即视频文件名
    
    # reduce_frame(video_file)

    # extract_i_frames("cut.mp4") #输入课程视频，输出命名好的key_frames文件夹
    # OCR_i_frames() #识别key_frames的内容到'output.txt'

    # api_key = "34e4250edb1cc63dfde749a7c83fbbc1.2RbubnzAdoUcBmxM"
    # AI_extract_questions(api_key) 

    # json2word() 
    # cleanup()




    